﻿namespace Fiddler
{
    using System;
    using System.Collections.Generic;
    using System.IO;
    using System.Net.Sockets;
    using System.Text;
    using System.Threading;

    public class WebSocket : ITunnel
    {
        private int _iMsgCount;
        private long _lngEgressByteCount;
        private long _lngIngressByteCount;
        private Session _mySession;
        private byte[] arrRequestBytes;
        private byte[] arrResponseBytes;
        private bool bIsOpen;
        private bool bParseMessages = FiddlerApplication.Prefs.GetBoolPref("fiddler.websocket.ParseMessages", true);
        private bool bShowRawBytes = FiddlerApplication.Prefs.GetBoolPref("fiddler.websocket.ShowRawBytes", false);
        public List<WebSocketMessage> listMessages;
        private ClientPipe oCP;
        private AutoResetEvent oKeepTunnelAlive;
        private ServerPipe oSP;
        private string sName = "Unknown";
        private MemoryStream strmClientBytes;
        private MemoryStream strmServerBytes;

        private WebSocket(Session oSess, ClientPipe oFrom, ServerPipe oTo)
        {
            this.sName = "WebSocket #" + oSess.id.ToString();
            this._mySession = oSess;
            this.oCP = oFrom;
            this.oSP = oTo;
            this._mySession.SetBitFlag(SessionFlags.IsWebSocketTunnel, true);
            if (oSess.oFlags.ContainsKey("x-no-parse"))
            {
                this.bParseMessages = false;
            }
        }

        private void _CleanupWebSocket()
        {
            this.bIsOpen = false;
            this.arrRequestBytes = (byte[]) (this.arrResponseBytes = null);
            this.strmServerBytes = null;
            this.strmClientBytes = null;
            if (this.oCP != null)
            {
                this.oCP.End();
            }
            if (this.oSP != null)
            {
                this.oSP.End();
            }
            this.oCP = null;
            this.oSP = null;
            if (this._mySession != null)
            {
                if ((this._mySession.oResponse != null) && (this._mySession.oResponse.headers != null))
                {
                    this._mySession.oResponse["EndTime"] = DateTime.Now.ToString("HH:mm:ss.fff");
                    this._mySession.oResponse["ReceivedBytes"] = this._lngIngressByteCount.ToString();
                    this._mySession.oResponse["SentBytes"] = this._lngEgressByteCount.ToString();
                }
                this._mySession.Timers.ServerDoneResponse = this._mySession.Timers.ClientBeginResponse = this._mySession.Timers.ClientDoneResponse = DateTime.Now;
                this._mySession = null;
            }
        }

        private void _ParseAndSendClientMessages()
        {
            foreach (WebSocketMessage message in _ParseMessagesFromStream(this, ref this.strmClientBytes, true, true))
            {
                lock (this.listMessages)
                {
                    this.listMessages.Add(message);
                }
                FiddlerApplication.oExtensions.DoOnWebSocketMessage(message);
                if (!message.WasAborted)
                {
                    byte[] arrData = message.ToByteArray();
                    this.oSP.BeginSend(arrData, 0, arrData.Length, SocketFlags.None, new AsyncCallback(this.OnSentToServer), message);
                }
            }
        }

        private void _ParseAndSendServerMessages()
        {
            foreach (WebSocketMessage message in _ParseMessagesFromStream(this, ref this.strmServerBytes, false, true))
            {
                lock (this.listMessages)
                {
                    this.listMessages.Add(message);
                }
                FiddlerApplication.oExtensions.DoOnWebSocketMessage(message);
                if (!message.WasAborted)
                {
                    byte[] arrData = message.ToByteArray();
                    this.oCP.BeginSend(arrData, 0, arrData.Length, SocketFlags.None, new AsyncCallback(this.OnSentToClient), message);
                }
            }
        }

        private static WebSocketMessage[] _ParseMessagesFromStream(WebSocket wsOwner, ref MemoryStream strmData, bool bIsOutbound, bool bTrimAfterParsing)
        {
            List<WebSocketMessage> list = new List<WebSocketMessage>();
            strmData.Position = 0;
            long position = 0;
            while ((strmData.Length - strmData.Position) > 2)
            {
                byte[] buffer = new byte[2];
                strmData.Read(buffer, 0, buffer.Length);
                ulong num2 = (ulong) (buffer[1] & 0x7f);
                if (num2 == 0x7e)
                {
                    if (strmData.Length < (strmData.Position + 2))
                    {
                        break;
                    }
                    byte[] buffer2 = new byte[2];
                    strmData.Read(buffer2, 0, buffer2.Length);
                    num2 = (ulong) ((buffer2[0] << 8) + buffer2[1]);
                }
                else if (num2 == 0x7f)
                {
                    if (strmData.Length < (strmData.Position + 8))
                    {
                        break;
                    }
                    byte[] buffer3 = new byte[8];
                    strmData.Read(buffer3, 0, buffer3.Length);
                    num2 = (ulong) ((((((((buffer3[0] << 0x18) + (buffer3[1] << 0x10)) + (buffer3[2] << 8)) + buffer3[3]) + (buffer3[4] << 0x18)) + (buffer3[5] << 0x10)) + (buffer3[6] << 8)) + buffer3[7]);
                }
                bool flag = 0x80 == (buffer[1] & 0x80);
                if ((ulong)strmData.Length < ((((ulong)strmData.Position) + (ulong)num2) + (flag ? (ulong)4 : (ulong)0)))
                {
                    break;
                }
                WebSocketMessage item = new WebSocketMessage(wsOwner, Interlocked.Increment(ref wsOwner._iMsgCount), bIsOutbound);
                item.AssignHeader(buffer[0]);
                if (flag)
                {
                    byte[] buffer4 = new byte[4];
                    strmData.Read(buffer4, 0, buffer4.Length);
                    item.MaskingKey = buffer4;
                }
                byte[] buffer5 = new byte[num2];
                strmData.Read(buffer5, 0, buffer5.Length);
                item.PayloadData = buffer5;
                list.Add(item);
                position = strmData.Position;
            }
            strmData.Position = position;
            if (bTrimAfterParsing)
            {
                byte[] buffer6 = new byte[strmData.Length - position];
                strmData.Read(buffer6, 0, buffer6.Length);
                strmData.Dispose();
                strmData = new MemoryStream();
                strmData.Write(buffer6, 0, buffer6.Length);
            }
            return list.ToArray();
        }

        private void _PushClientBuffer(int iReadCount)
        {
            FiddlerApplication.Log.LogFormat("{0} Pushing {1} bytes from client WebSocket", new object[] { this, iReadCount });
            this.strmClientBytes.Write(this.arrRequestBytes, 0, iReadCount);
            this._ParseAndSendClientMessages();
        }

        private void _PushServerBuffer(int iReadCount)
        {
            FiddlerApplication.Log.LogFormat("{0} - Pushing {1} bytes from server WebSocket\n{2}", new object[] { this, iReadCount, Utilities.ByteArrayToHexView(this.arrResponseBytes, 0x10, iReadCount) });
            this.strmServerBytes.Write(this.arrResponseBytes, 0, iReadCount);
            this._ParseAndSendServerMessages();
        }

        private string[] _ReadHeadersFromStream(Stream oFS)
        {
            List<byte> list = new List<byte>();
            bool flag = false;
            bool flag2 = true;
            for (int i = oFS.ReadByte(); -1 != i; i = oFS.ReadByte())
            {
                if (i == 13)
                {
                    flag = true;
                }
                else if (flag && (i == 10))
                {
                    if (flag2)
                    {
                        break;
                    }
                    flag2 = true;
                    list.Add(13);
                    list.Add(10);
                }
                else
                {
                    flag = flag2 = false;
                    list.Add((byte) i);
                }
            }
            return Encoding.ASCII.GetString(list.ToArray()).Split(new string[] { "\r\n" }, StringSplitOptions.RemoveEmptyEntries);
        }

        public void CloseTunnel()
        {
            if (this.bParseMessages)
            {
                FiddlerApplication.Log.LogString("Close WebSocket Tunnel: " + Environment.StackTrace);
            }
            try
            {
                if (this.oKeepTunnelAlive != null)
                {
                    this.oKeepTunnelAlive.Set();
                }
                else
                {
                    this._CleanupWebSocket();
                }
            }
            catch (Exception exception)
            {
                FiddlerApplication.Log.LogString("Error closing oKeepTunnelAlive. " + exception.Message + "\n" + exception.StackTrace);
            }
        }

        internal static void CreateTunnel(Session oSession)
        {
            if ((((oSession != null) && (oSession.oRequest != null)) && ((oSession.oRequest.headers != null) && (oSession.oRequest.pipeClient != null))) && ((oSession.oResponse != null) && (oSession.oResponse.pipeServer != null)))
            {
                FiddlerApplication.Log.LogFormat("WebSocket Key Verification\nActual:\t{0}\nFiddler:\t{1}", new object[] { oSession.oResponse.headers["Sec-WebSocket-Accept"], WebSocketUtility.ComputeAcceptKey(oSession.oRequest.headers["Sec-WebSocket-Key"]) });
                ClientPipe pipeClient = oSession.oRequest.pipeClient;
                oSession.oRequest.pipeClient = null;
                ServerPipe pipeServer = oSession.oResponse.pipeServer;
                oSession.oResponse.pipeServer = null;
                WebSocket socket = new WebSocket(oSession, pipeClient, pipeServer);
                oSession.__oTunnel = socket;
                Thread thread = new Thread(new ThreadStart(socket.RunTunnel));
                thread.IsBackground = true;
                thread.Start();
            }
        }

        internal static void LoadWebSocketMessagesFromStream(Session oS, Stream strmWSMessages)
        {
            try
            {
                WebSocket socket = new WebSocket(oS, null, null);
                socket.sName = string.Format("SAZ-Session#{0}", oS.id);
                oS.__oTunnel = socket;
                socket.ReadWebSocketMessageListFromStream(strmWSMessages);
            }
            finally
            {
                strmWSMessages.Dispose();
            }
        }

        protected void OnReceiveFromClient(IAsyncResult ar)
        {
            try
            {
                int iReadCount = this.oCP.EndReceive(ar);
                if (iReadCount > 0)
                {
                    this._lngEgressByteCount += iReadCount;
                    if (this.bParseMessages)
                    {
                        this._PushClientBuffer(iReadCount);
                    }
                    else
                    {
                        this.oSP.BeginSend(this.arrRequestBytes, 0, iReadCount, SocketFlags.None, new AsyncCallback(this.OnSentToServer), this.oSP);
                    }
                    this.oCP.BeginReceive(this.arrRequestBytes, 0, this.arrRequestBytes.Length, SocketFlags.None, new AsyncCallback(this.OnReceiveFromClient), this.oCP);
                }
                else
                {
                    if (this.bParseMessages)
                    {
                        FiddlerApplication.Log.LogFormat("[{0}] Read from Client returned error: {1}", new object[] { this.sName, iReadCount });
                    }
                    this.CloseTunnel();
                }
            }
            catch (Exception exception)
            {
                if (this.bParseMessages)
                {
                    FiddlerApplication.Log.LogFormat("[{0}] Read from Client failed... {1}", new object[] { this.sName, exception.Message });
                }
                this.CloseTunnel();
            }
        }

        protected void OnReceiveFromServer(IAsyncResult ar)
        {
            try
            {
                int iReadCount = this.oSP.EndReceive(ar);
                if (iReadCount > 0)
                {
                    this._lngIngressByteCount += iReadCount;
                    if (this.bParseMessages)
                    {
                        this._PushServerBuffer(iReadCount);
                    }
                    else
                    {
                        this.oCP.BeginSend(this.arrResponseBytes, 0, iReadCount, SocketFlags.None, new AsyncCallback(this.OnSentToClient), this.oCP);
                    }
                    this.oSP.BeginReceive(this.arrResponseBytes, 0, this.arrResponseBytes.Length, SocketFlags.None, new AsyncCallback(this.OnReceiveFromServer), this.oSP);
                }
                else
                {
                    if (this.bParseMessages)
                    {
                        FiddlerApplication.Log.LogFormat("[{0}] Read from Server returned error: {1}", new object[] { this.sName, iReadCount });
                    }
                    this.CloseTunnel();
                }
            }
            catch (Exception exception)
            {
                if (this.bParseMessages)
                {
                    FiddlerApplication.Log.LogFormat("[{0}] Read from Server failed... {1}", new object[] { this.sName, exception.Message });
                }
                this.CloseTunnel();
            }
        }

        protected void OnSentToClient(IAsyncResult ar)
        {
            try
            {
                this.oCP.EndSend(ar);
            }
            catch (Exception exception)
            {
                if (this.bParseMessages)
                {
                    FiddlerApplication.Log.LogFormat("[{0}] WebSocket OnSentToClient failed... {1}", new object[] { this.sName, exception.Message });
                }
                this.CloseTunnel();
            }
        }

        protected void OnSentToServer(IAsyncResult ar)
        {
            try
            {
                this.oSP.EndSend(ar);
            }
            catch (Exception exception)
            {
                if (this.bParseMessages)
                {
                    FiddlerApplication.Log.LogFormat("[{0}] OnSendToServer failed... {1}", new object[] { this.sName, exception.Message });
                }
                this.CloseTunnel();
            }
        }

        internal bool ReadWebSocketMessageListFromStream(Stream oFS)
        {
            try
            {
                string[] strArray = this._ReadHeadersFromStream(oFS);
                List<WebSocketMessage> list = new List<WebSocketMessage>();
                strArray = this._ReadHeadersFromStream(oFS);
                while (strArray.Length > 0)
                {
                    int count = 0;
                    bool bIsOutbound = false;
                    foreach (string str in strArray)
                    {
                        if (str.StartsWith("Request-Length:"))
                        {
                            bIsOutbound = true;
                            count = int.Parse(str.Substring(0x10));
                            break;
                        }
                        if (str.StartsWith("Response-Length:"))
                        {
                            bIsOutbound = false;
                            count = int.Parse(str.Substring(0x11));
                            break;
                        }
                    }
                    if (count < 1)
                    {
                        throw new InvalidDataException("Missing size indication.");
                    }
                    byte[] buffer = new byte[count];
                    oFS.Read(buffer, 0, count);
                    MemoryStream strmData = new MemoryStream(buffer);
                    WebSocketMessage[] messageArray = _ParseMessagesFromStream(this, ref strmData, bIsOutbound, false);
                    list.Add(messageArray[0]);
                    if ((-1 == oFS.ReadByte()) || (-1 == oFS.ReadByte()))
                    {
                        strArray = new string[0];
                    }
                    else
                    {
                        strArray = this._ReadHeadersFromStream(oFS);
                    }
                }
                this.listMessages = list;
                return true;
            }
            catch (Exception)
            {
                return false;
            }
        }

        private void RunTunnel()
        {
            if (FiddlerApplication.oProxy != null)
            {
                this.arrRequestBytes = new byte[0x4000];
                this.arrResponseBytes = new byte[0x4000];
                if (this.bParseMessages)
                {
                    this.strmClientBytes = new MemoryStream();
                    this.strmServerBytes = new MemoryStream();
                    this.listMessages = new List<WebSocketMessage>();
                }
                this.bIsOpen = true;
                try
                {
                    this.oCP.BeginReceive(this.arrRequestBytes, 0, this.arrRequestBytes.Length, SocketFlags.None, new AsyncCallback(this.OnReceiveFromClient), this.oCP);
                    this.oSP.BeginReceive(this.arrResponseBytes, 0, this.arrResponseBytes.Length, SocketFlags.None, new AsyncCallback(this.OnReceiveFromServer), this.oSP);
                    FiddlerApplication.UIInvokeAsync(new updateUIDelegate(FiddlerApplication._frmMain.finishSession), new object[] { this._mySession });
                    this.WaitForCompletion();
                }
                catch (Exception)
                {
                }
                this.CloseTunnel();
            }
        }

        public override string ToString()
        {
            return string.Format("Session{0}.WebSocket'{1}'", (this._mySession == null) ? -1 : this._mySession.id, this.sName);
        }

        private void WaitForCompletion()
        {
            AutoResetEvent oKeepTunnelAlive = this.oKeepTunnelAlive;
            this.oKeepTunnelAlive = new AutoResetEvent(false);
            this.oKeepTunnelAlive.WaitOne();
            this.oKeepTunnelAlive.Close();
            this.oKeepTunnelAlive = null;
        }

        internal bool WriteWebSocketMessageListToStream(Stream oFS)
        {
            oFS.WriteByte(13);
            oFS.WriteByte(10);
            lock (this.listMessages)
            {
                foreach (WebSocketMessage message in this.listMessages)
                {
                    byte[] buffer = message.ToByteArray();
                    string s = string.Format("{0}: {1}\r\nID: {2}\r\n\r\n", message.IsOutbound ? "Request-Length" : "Response-Length", buffer.Length, message.ID);
                    byte[] bytes = Encoding.ASCII.GetBytes(s);
                    oFS.Write(bytes, 0, bytes.Length);
                    oFS.Write(buffer, 0, buffer.Length);
                    oFS.WriteByte(13);
                    oFS.WriteByte(10);
                }
            }
            return true;
        }

        public long EgressByteCount
        {
            get
            {
                return this._lngEgressByteCount;
            }
        }

        public long IngressByteCount
        {
            get
            {
                return this._lngIngressByteCount;
            }
        }

        internal bool IsBlind
        {
            get
            {
                return !this.bParseMessages;
            }
            set
            {
                this.bParseMessages = !value;
            }
        }

        public bool IsOpen
        {
            get
            {
                return this.bIsOpen;
            }
        }
    }
}

